package org.duanyu.helper;

import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.FrameGrabber;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.model.Media;
import org.springframework.core.io.PathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 视频转图片
 */
public class FrameHelper {
    private static final Logger logger = LoggerFactory.getLogger(FrameHelper.class);

    public static List<Media> load(MultipartFile multipartFile) {
        List<Media> mediaList = new ArrayList<>();
        // 设置图片输出路径（可以考虑配置化）
        String imagePath = "C:\\Users\\29712\\Desktop\\pic\\output\\";
        try {
            // 初始化视频抓取工具
            FFmpegFrameGrabber ff = new FFmpegFrameGrabber(multipartFile.getInputStream());
            // 帧到图像转换工具
            Java2DFrameConverter converter = new Java2DFrameConverter();
            // 抓取内容
            ff.start();
            ff.setFormat("mp4");

            // 最大帧数限制（防止无限循环）
            int maxFrames = 1000; // 可以根据实际情况调整
            Frame frame;
            int frameCount = 0;
            // 遍历视频帧，每一帧都转为图片保存
            while ((frame = ff.grabFrame()) != null && frameCount < maxFrames) {
                // 判断空帧
                if (frame.image == null) {
                    continue;
                }

                // 将视频帧转为BufferedImage
                BufferedImage image = converter.getBufferedImage(frame);

                // 图片路径
                String imageFilePath = imagePath + frameCount + ".png";
                File file = new File(imageFilePath);
                // 将图像写入文件，并记录图像的路径
                ImageIO.write(image, "png", file);
                Media media = new Media(
                        MimeTypeUtils.IMAGE_PNG,
                        new PathResource(imageFilePath));
                mediaList.add(media);

                frameCount++;
            }
            ff.stop();
            // 检查图像数量是否在有效范围内
            if (mediaList.size() < 4) {
                logger.warn("提取的图像数量不足4张，当前数量：{}", mediaList.size());
            } else if (mediaList.size() > 512) {
                logger.warn("提取的图像数量超过512张，当前数量：{}", mediaList.size());
                // 可以考虑截断到512张
                mediaList = mediaList.subList(0, 512);
            }

        } catch (FFmpegFrameGrabber.Exception e) {
            logger.error("FFmpegFrameGrabber异常", e);
            throw new RuntimeException(e);
        } catch (FrameGrabber.Exception e) {
            logger.error("FrameGrabber异常", e);
            throw new RuntimeException(e);
        } catch (IOException e) {
            logger.error("IO异常", e);
            throw new RuntimeException(e);
        }

        return mediaList;
    }
}